import sys
import cv2
import numpy as np
from PIL import Image, ImageTk
from PyQt5 import QtCore, QtGui, QtWidgets


class ImageProcessor:
    """图像处理核心功能类"""

    def __init__(self):
        self.image = None
        self.processed_image = None
        self.geometry = None  # 形状识别结果
        self.faces = None  # 人脸检测结果

    # 1. 基础图像处理功能
    def load_image(self, path):
        """加载图像"""
        self.image = cv2.imread(path)
        self.processed_image = self.image.copy()
        return self.convert_to_qt_image(self.processed_image)

    def convert_to_qt_image(self, img):
        """转换为Qt显示格式"""
        if img is not None:
            rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(rgb_image)
            return ImageTk.PhotoImage(pil_image)
        return None

    def rotate_image(self, angle):
        """旋转图像"""
        rows, cols = self.processed_image.shape[:2]
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
        self.processed_image = cv2.warpAffine(self.processed_image, M, (cols, rows))

    def scale_image(self, scale_factor):
        """缩放图像"""
        width = int(self.processed_image.shape[1] * scale_factor)
        height = int(self.processed_image.shape[0] * scale_factor)
        self.processed_image = cv2.resize(self.processed_image, (width, height))

    def crop_image(self, x1, y1, x2, y2):
        """裁剪图像"""
        self.processed_image = self.processed_image[y1:y2, x1:x2]

    def adjust_color(self, brightness, contrast, saturation):
        """调整色彩"""
        hsv = cv2.cvtColor(self.processed_image, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)

        # 亮度调整（v通道）
        adjusted_v = cv2.addWeighted(v, contrast, v, 0, brightness)
        # 饱和度调整（s通道）
        adjusted_s = np.clip(s * saturation, 0, 255).astype(np.uint8)

        adjusted_hsv = cv2.merge((h, adjusted_s, adjusted_v))
        self.processed_image = cv2.cvtColor(adjusted_hsv, cv2.COLOR_HSV2BGR)

    def smooth_image(self, kernel_size=5):
        """图像平滑处理"""
        self.processed_image = cv2.GaussianBlur(self.processed_image, (kernel_size, kernel_size), 0)

    def detect_shapes(self):
        """形状识别"""
        if self.processed_image is None:
            return None

        gray = cv2.cvtColor(self.processed_image, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        shapes = []
        for cnt in contours:
            approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), True)
            if len(approx) == 3:
                shapes.append(('Triangle', approx))
            elif len(approx) == 4:
                x, y, w, h = cv2.boundingRect(cnt)
                ar = w / float(h)
                shapes.append(('Rectangle' if abs(ar - 1) > 0.1 else 'Square', approx))
            elif len(approx) == 5:
                shapes.append(('Pentagon', approx))
            elif len(approx) == 6:
                shapes.append(('Hexagon', approx))
            else:
                (x, y), radius = cv2.minEnclosingCircle(cnt)
                shapes.append(('Circle', np.int0([[x, y, radius]])))

        self.geometry = shapes
        return self.draw_shape_contours()

    def draw_shape_contours(self):
        """绘制形状轮廓"""
        if self.geometry:
            img = self.processed_image.copy()
            for shape in self.geometry:
                name, approx = shape
                if name == 'Circle':
                    (x, y, r) = approx[0]
                    cv2.circle(img, (int(x), int(y)), int(r), (0, 255, 0), 2)
                    cv2.putText(img, f'Circle', (int(x) - 20, int(y) - 20),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                else:
                    cv2.drawContours(img, [approx], -1, (0, 255, 0), 2)
                    x, y = approx[0][0]
                    cv2.putText(img, name, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            self.processed_image = img
        return self.convert_to_qt_image(self.processed_image)

    # 2. 计算机视觉算法（简化版人脸检测，使用Haar特征）
    def detect_faces(self):
        """人脸检测（使用OpenCV Haar级联分类器）"""
        if self.processed_image is None:
            return None

        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        gray = cv2.cvtColor(self.processed_image, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)

        for (x, y, w, h) in faces:
            cv2.rectangle(self.processed_image, (x, y), (x + w, y + h), (255, 0, 0), 2)
            cv2.putText(self.processed_image, 'Face', (x, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
        self.faces = faces
        return self.convert_to_qt_image(self.processed_image)

    # 3. 视频处理功能
    def process_video_frame(self, frame):
        """实时视频处理"""
        if frame is not None:
            self.processed_image = frame.copy()
            # 应用形状识别和人脸检测
            self.detect_shapes()
            self.detect_faces()
            return self.convert_to_qt_image(self.processed_image)
        return None


class MainWindow(QtWidgets.QMainWindow):
    """主窗口界面类"""

    def __init__(self):
        super().__init__()
        self.processor = ImageProcessor()
        self.initUI()
        self.video_cap = None  # 视频流对象
        self.timer = None  # 视频定时器

    def initUI(self):
        # 窗口设置
        self.setWindowTitle("PyQt5图像处理平台")
        self.setGeometry(100, 100, 1000, 800)

        # 中央部件
        central_widget = QtWidgets.QWidget()
        self.setCentralWidget(central_widget)
        layout = QtWidgets.QVBoxLayout(central_widget)

        # 菜单栏
        self.create_menu()

        # 工具栏
        self.create_toolbar()

        # 图像/视频显示区域
        self.image_label = QtWidgets.QLabel()
        self.image_label.setMinimumSize(640, 480)
        self.image_label.setAlignment(QtCore.Qt.AlignCenter)
        self.image_label.setText("请加载图像或启动视频")
        layout.addWidget(self.image_label)

        # 参数设置区域
        param_layout = QtWidgets.QHBoxLayout()

        # 色彩调整滑块
        color_widget = QtWidgets.QGroupBox("色彩调整")
        color_layout = QtWidgets.QVBoxLayout(color_widget)

        brightness_layout = QtWidgets.QHBoxLayout()
        brightness_layout.addWidget(QtWidgets.QLabel("亮度"))
        self.brightness_slider = QtWidgets.QSlider(QtCore.Qt.Horizontal)
        self.brightness_slider.setRange(-100, 100)
        self.brightness_slider.setValue(0)
        self.brightness_slider.valueChanged.connect(self.adjust_color_params)
        brightness_layout.addWidget(self.brightness_slider)
        self.brightness_value = QtWidgets.QLabel("0")
        self.brightness_value.setFixedWidth(30)
        brightness_layout.addWidget(self.brightness_value)
        color_layout.addLayout(brightness_layout)

        contrast_layout = QtWidgets.QHBoxLayout()
        contrast_layout.addWidget(QtWidgets.QLabel("对比度"))
        self.contrast_slider = QtWidgets.QSlider(QtCore.Qt.Horizontal)
        self.contrast_slider.setRange(0, 200)
        self.contrast_slider.setValue(100)
        self.contrast_slider.valueChanged.connect(self.adjust_color_params)
        contrast_layout.addWidget(self.contrast_slider)
        self.contrast_value = QtWidgets.QLabel("1.0")
        self.contrast_value.setFixedWidth(30)
        contrast_layout.addWidget(self.contrast_value)
        color_layout.addLayout(contrast_layout)

        saturation_layout = QtWidgets.QHBoxLayout()
        saturation_layout.addWidget(QtWidgets.QLabel("饱和度"))
        self.saturation_slider = QtWidgets.QSlider(QtCore.Qt.Horizontal)
        self.saturation_slider.setRange(0, 200)
        self.saturation_slider.setValue(100)
        self.saturation_slider.valueChanged.connect(self.adjust_color_params)
        saturation_layout.addWidget(self.saturation_slider)
        self.saturation_value = QtWidgets.QLabel("1.0")
        self.saturation_value.setFixedWidth(30)
        saturation_layout.addWidget(self.saturation_value)
        color_layout.addLayout(saturation_layout)

        param_layout.addWidget(color_widget)

        # 平滑处理控件
        smooth_widget = QtWidgets.QGroupBox("平滑处理")
        smooth_layout = QtWidgets.QVBoxLayout(smooth_widget)

        kernel_layout = QtWidgets.QHBoxLayout()
        kernel_layout.addWidget(QtWidgets.QLabel("核大小"))
        self.kernel_size_spin = QtWidgets.QSpinBox()
        self.kernel_size_spin.setRange(1, 21)
        self.kernel_size_spin.setValue(5)
        self.kernel_size_spin.setSingleStep(2)  # 只能选择奇数
        self.kernel_size_spin.valueChanged.connect(self.apply_smoothing)
        kernel_layout.addWidget(self.kernel_size_spin)
        smooth_layout.addLayout(kernel_layout)

        smooth_btn = QtWidgets.QPushButton("应用平滑")
        smooth_btn.clicked.connect(self.apply_smoothing)
        smooth_layout.addWidget(smooth_btn)

        param_layout.addWidget(smooth_widget)

        # 形状识别控件
        shape_widget = QtWidgets.QGroupBox("形状识别")
        shape_layout = QtWidgets.QVBoxLayout(shape_widget)

        shape_btn = QtWidgets.QPushButton("检测形状")
        shape_btn.clicked.connect(self.detect_shapes)
        shape_layout.addWidget(shape_btn)

        param_layout.addWidget(shape_widget)

        # 人脸检测控件
        face_widget = QtWidgets.QGroupBox("人脸检测")
        face_layout = QtWidgets.QVBoxLayout(face_widget)

        face_btn = QtWidgets.QPushButton("检测人脸")
        face_btn.clicked.connect(self.detect_faces)
        face_layout.addWidget(face_btn)

        param_layout.addWidget(face_widget)

        layout.addLayout(param_layout)

        # 状态栏
        self.statusBar().showMessage("就绪")

    def create_menu(self):
        menu_bar = self.menuBar()

        # 文件菜单
        file_menu = menu_bar.addMenu("文件")
        load_action = QtWidgets.QAction("加载图像", self)
        load_action.setShortcut("Ctrl+O")
        load_action.triggered.connect(self.load_image_dialog)
        file_menu.addAction(load_action)

        save_action = QtWidgets.QAction("保存图像", self)
        save_action.setShortcut("Ctrl+S")
        save_action.triggered.connect(self.save_image)
        file_menu.addAction(save_action)

        file_menu.addSeparator()

        exit_action = QtWidgets.QAction("退出", self)
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)

        # 编辑菜单
        edit_menu = menu_bar.addMenu("编辑")
        rotate_action = QtWidgets.QAction("旋转图像", self)
        rotate_action.setShortcut("Ctrl+R")
        rotate_action.triggered.connect(self.rotate_image_dialog)
        edit_menu.addAction(rotate_action)

        scale_action = QtWidgets.QAction("缩放图像", self)
        scale_action.setShortcut("Ctrl+Z")
        scale_action.triggered.connect(self.scale_image_dialog)
        edit_menu.addAction(scale_action)

        crop_action = QtWidgets.QAction("裁剪图像", self)
        crop_action.setShortcut("Ctrl+C")
        crop_action.triggered.connect(self.crop_image_dialog)
        edit_menu.addAction(crop_action)

        # 处理菜单
        process_menu = menu_bar.addMenu("处理")
        shape_action = QtWidgets.QAction("形状识别", self)
        shape_action.triggered.connect(self.detect_shapes)
        process_menu.addAction(shape_action)

        face_action = QtWidgets.QAction("人脸检测", self)
        face_action.triggered.connect(self.detect_faces)
        process_menu.addAction(face_action)

        # 视频菜单
        video_menu = menu_bar.addMenu("视频")
        play_video_action = QtWidgets.QAction("播放视频", self)
        play_video_action.triggered.connect(self.play_video_dialog)
        video_menu.addAction(play_video_action)

        live_camera_action = QtWidgets.QAction("实时摄像头", self)
        live_camera_action.triggered.connect(self.start_live_camera)
        video_menu.addAction(live_camera_action)

        # 帮助菜单
        help_menu = menu_bar.addMenu("帮助")
        about_action = QtWidgets.QAction("关于", self)
        about_action.triggered.connect(self.show_about)
        help_menu.addAction(about_action)

    def create_toolbar(self):
        toolbar = self.addToolBar("工具")
        toolbar.setMovable(False)

        load_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("document-open"), "加载图像", self)
        load_btn.triggered.connect(self.load_image_dialog)
        toolbar.addAction(load_btn)

        save_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("document-save"), "保存图像", self)
        save_btn.triggered.connect(self.save_image)
        toolbar.addAction(save_btn)

        toolbar.addSeparator()

        rotate_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("object-rotate-right"), "旋转", self)
        rotate_btn.triggered.connect(self.rotate_image_dialog)
        toolbar.addAction(rotate_btn)

        scale_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("zoom-in"), "缩放", self)
        scale_btn.triggered.connect(self.scale_image_dialog)
        toolbar.addAction(scale_btn)

        crop_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("transform-crop"), "裁剪", self)
        crop_btn.triggered.connect(self.crop_image_dialog)
        toolbar.addAction(crop_btn)

        toolbar.addSeparator()

        shape_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("draw-polygon"), "形状识别", self)
        shape_btn.triggered.connect(self.detect_shapes)
        toolbar.addAction(shape_btn)

        face_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("face-smile"), "人脸检测", self)
        face_btn.triggered.connect(self.detect_faces)
        toolbar.addAction(face_btn)

        toolbar.addSeparator()

        video_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("media-playback-start"), "播放视频", self)
        video_btn.triggered.connect(self.play_video_dialog)
        toolbar.addAction(video_btn)

        camera_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("camera-photo"), "摄像头", self)
        camera_btn.triggered.connect(self.start_live_camera)
        toolbar.addAction(camera_btn)

    # ---------- 图像操作 ----------
    def load_image_dialog(self):
        path, _ = QtWidgets.QFileDialog.getOpenFileName(
            self, "加载图像", "",
            "图像文件 (*.jpg *.jpeg *.png *.bmp *.gif);;所有文件 (*)"
        )
        if path:
            self.stop_video()  # 停止任何正在运行的视频
            qt_image = self.processor.load_image(path)
            if qt_image:
                self.image_label.setPixmap(QtGui.QPixmap.fromImage(qt_image.toqpixmap()))
                self.statusBar().showMessage(f"已加载图像：{path}")

    def save_image(self):
        if self.processor.processed_image is not None:
            path, _ = QtWidgets.QFileDialog.getSaveFileName(
                self, "保存图像", "",
                "PNG图像 (*.png);;JPEG图像 (*.jpg);;BMP图像 (*.bmp);;所有文件 (*)"
            )
            if path:
                cv2.imwrite(path, self.processor.processed_image)
                self.statusBar().showMessage(f"图像已保存至：{path}")

    def rotate_image_dialog(self):
        if self.processor.processed_image is not None:
            angle, ok = QtWidgets.QInputDialog.getDouble(
                self, "旋转图像", "输入角度（度）：",
                0, -360, 360, 1
            )
            if ok:
                self.processor.rotate_image(angle)
                self.update_display()

    def scale_image_dialog(self):
        if self.processor.processed_image is not None:
            factor, ok = QtWidgets.QInputDialog.getDouble(
                self, "缩放图像", "输入缩放比例：",
                1, 0.1, 10, 2
            )
            if ok:
                self.processor.scale_image(factor)
                self.update_display()

    def crop_image_dialog(self):
        if self.processor.processed_image is not None:
            h, w = self.processor.processed_image.shape[:2]

            dialog = QtWidgets.QDialog(self)
            dialog.setWindowTitle("裁剪图像")
            dialog.setMinimumWidth(300)

            layout = QtWidgets.QFormLayout(dialog)

            x1 = QtWidgets.QSpinBox()
            x1.setRange(0, w - 1)
            y1 = QtWidgets.QSpinBox()
            y1.setRange(0, h - 1)
            x2 = QtWidgets.QSpinBox()
            x2.setRange(1, w)
            x2.setValue(w)
            y2 = QtWidgets.QSpinBox()
            y2.setRange(1, h)
            y2.setValue(h)

            layout.addRow("左边界(X1):", x1)
            layout.addRow("上边界(Y1):", y1)
            layout.addRow("右边界(X2):", x2)
            layout.addRow("下边界(Y2):", y2)

            btn_box = QtWidgets.QDialogButtonBox(
                QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel
            )
            layout.addWidget(btn_box)

            btn_box.accepted.connect(lambda: self.processor.crop_image(x1.value(), y1.value(), x2.value(), y2.value()))
            btn_box.rejected.connect(dialog.reject)

            if dialog.exec_() == QtWidgets.QDialog.Accepted:
                self.update_display()

    def adjust_color_params(self):
        if self.processor.processed_image is not None:
            brightness = self.brightness_slider.value()
            contrast = self.contrast_slider.value() / 100.0  # 转换为0.0-2.0范围
            saturation = self.saturation_slider.value() / 100.0  # 转换为0.0-2.0范围

            # 更新显示值
            self.brightness_value.setText(str(brightness))
            self.contrast_value.setText(f"{contrast:.1f}")
            self.saturation_value.setText(f"{saturation:.1f}")

            self.processor.adjust_color(brightness, contrast, saturation)
            self.update_display()

    def apply_smoothing(self):
        if self.processor.processed_image is not None:
            kernel_size = self.kernel_size_spin.value()
            if kernel_size % 2 == 0:  # 确保核大小为奇数
                kernel_size += 1
                self.kernel_size_spin.setValue(kernel_size)

            self.processor.smooth_image(kernel_size)
            self.update_display()

    def detect_shapes(self):
        if self.processor.processed_image is not None:
            qt_image = self.processor.detect_shapes()
            if qt_image:
                self.image_label.setPixmap(QtGui.QPixmap.fromImage(qt_image.toqpixmap()))
                if self.processor.geometry:
                    count = len(self.processor.geometry)
                    self.statusBar().showMessage(f"检测到 {count} 个形状")
                else:
                    self.statusBar().showMessage("未检测到形状")

    def detect_faces(self):
        if self.processor.processed_image is not None:
            qt_image = self.processor.detect_faces()
            if qt_image:
                self.image_label.setPixmap(QtGui.QPixmap.fromImage(qt_image.toqpixmap()))
                if self.processor.faces is not None and len(self.processor.faces) > 0:
                    count = len(self.processor.faces)
                    self.statusBar().showMessage(f"检测到 {count} 个人脸")
                else:
                    self.statusBar().showMessage("未检测到人脸")

    # ---------- 视频处理 ----------
    def play_video_dialog(self):
        path, _ = QtWidgets.QFileDialog.getOpenFileName(
            self, "播放视频", "",
            "视频文件 (*.mp4 *.avi *.mkv *.mov *.wmv);;所有文件 (*)"
        )
        if path:
            self.start_video(path)

    def start_live_camera(self):
        self.start_video(0)  # 0表示默认摄像头

    def start_video(self, source):
        self.stop_video()  # 停止任何正在运行的视频

        self.video_cap = cv2.VideoCapture(source)
        if not self.video_cap.isOpened():
            self.statusBar().showMessage("无法打开视频源")
            return

        self.timer = QtCore.QTimer(self)
        self.timer.timeout.connect(self.update_video_frame)
        self.timer.start(30)  # 约30fps

        if isinstance(source, str):
            self.statusBar().showMessage(f"正在播放视频: {source}")
        else:
            self.statusBar().showMessage("正在捕获摄像头")

    def stop_video(self):
        if self.timer and self.timer.isActive():
            self.timer.stop()

        if self.video_cap:
            self.video_cap.release()
            self.video_cap = None

    def update_video_frame(self):
        if self.video_cap and self.video_cap.isOpened():
            ret, frame = self.video_cap.read()
            if ret:
                qt_image = self.processor.process_video_frame(frame)
                if qt_image:
                    self.image_label.setPixmap(QtGui.QPixmap.fromImage(qt_image.toqpixmap()))
            else:
                self.stop_video()
                self.statusBar().showMessage("视频播放结束")

    # ---------- 辅助方法 ----------
    def update_display(self):
        """更新显示区域"""
        if self.processor.processed_image is not None:
            qt_image = self.processor.convert_to_qt_image(self.processor.processed_image)
            if qt_image:
                self.image_label.setPixmap(QtGui.QPixmap.fromImage(qt_image.toqpixmap()))

    def show_about(self):
        """显示关于对话框"""
        QtWidgets.QMessageBox.about(
            self, "关于图像处理平台",
            "PyQt5图像处理平台\n\n"
            "版本: 1.0\n"
            "功能: 图像编辑、形状识别、人脸检测、视频处理\n\n"
            "© 2025 图像处理平台开发团队"
        )

    def closeEvent(self, event):
        """关闭窗口时停止所有视频处理"""
        self.stop_video()
        event.accept()


if __name__ == "__main__":
    # 确保中文显示正常
    import matplotlib

    matplotlib.use('Agg')
    matplotlib.rcParams['font.family'] = ['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC']

    app = QtWidgets.QApplication(sys.argv)
    app.setStyle('Fusion')  # 使用Fusion风格，跨平台一致性更好

    # 设置应用程序样式
    palette = QtGui.QPalette()
    palette.setColor(QtGui.QPalette.Window, QtGui.QColor(53, 53, 53))
    palette.setColor(QtGui.QPalette.WindowText, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.Base, QtGui.QColor(25, 25, 25))
    palette.setColor(QtGui.QPalette.AlternateBase, QtGui.QColor(53, 53, 53))
    palette.setColor(QtGui.QPalette.ToolTipBase, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.ToolTipText, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.Text, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.Button, QtGui.QColor(53, 53, 53))
    palette.setColor(QtGui.QPalette.ButtonText, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.BrightText, QtCore.Qt.red)
    palette.setColor(QtGui.QPalette.Link, QtGui.QColor(42, 130, 218))
    palette.setColor(QtGui.QPalette.Highlight, QtGui.QColor(42, 130, 218))
    palette.setColor(QtGui.QPalette.HighlightedText, QtCore.Qt.black)
    app.setPalette(palette)

    window = MainWindow()
    window.show()
    sys.exit(app.exec_())